Skip to content

LCORE-1350: Increased type-safety of tool utilities#1229

Merged
tisnik merged 1 commit intolightspeed-core:mainfrom
asimurka:post_responses_impl
Feb 26, 2026
Merged

LCORE-1350: Increased type-safety of tool utilities#1229
tisnik merged 1 commit intolightspeed-core:mainfrom
asimurka:post_responses_impl

Conversation

@asimurka
Copy link
Contributor

@asimurka asimurka commented Feb 26, 2026

Description

This PR refactors:

  • utility functions for getting available tools for request to use LLS API model types instead of plain dicts to make them more type-safe.
  • extends ResponsesApiParams model to contain all currently available attributes supported by LLS
  • makes query requests explicitly drop attributes that are not set (None)
  • introduces optional attributes notation

Type of change

  • Refactor
  • New feature
  • Bug fix
  • CVE fix
  • Optimization
  • Documentation Update
  • Configuration Update
  • Bump-up service version
  • Bump-up dependent library
  • Bump-up library or tool used for development (does not change the final image)
  • CI configuration change
  • Konflux configuration change
  • Unit tests improvement
  • Integration tests improvement
  • End to end tests improvement
  • Benchmarks improvement

Tools used to create PR

Identify any AI code assistants used in this PR (for transparency and review context)

  • Assisted-by: Cursor

Related Tickets & Documents

Checklist before requesting a review

  • I have performed a self-review of my code.
  • PR has passed all pre-merge test jobs.
  • If it is a core feature, I have added thorough tests.

Testing

  • Please provide detailed steps to perform tests related to this code change.
  • How were the fix/results from this change verified? Please provide relevant screenshots or results.

Summary by CodeRabbit

  • New Features

    • Expanded response configuration with new parameters: conversation history, response inclusion options, inference iterations, tool call limits, metadata, parallel tool execution, previous response references, custom prompts, temperature control, and tool choice strategies.
    • Enhanced input flexibility supporting both text and rich structured response content.
  • Bug Fixes

    • Improved payload cleanliness by excluding null fields from API responses.
    • Refined input type handling with explicit string casting in moderation and persistence workflows.
  • Refactor

    • Standardized tool definitions to use structured types instead of generic dictionaries for better type safety and consistency.

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Feb 26, 2026

Walkthrough

This PR refactors the type system for response handling and tools management. It introduces structured type definitions (InputTool, InputToolFileSearch, InputToolMCP) to replace dict-based tool representations, expands ResponsesApiParams with new fields for enhanced response composition, standardizes Optional type annotations throughout responses utilities, and adds explicit string casting and None-field filtering in endpoint response handlers.

Changes

Cohort / File(s) Summary
Query Endpoints
src/app/endpoints/query.py, src/app/endpoints/streaming_query.py
Added explicit str casting for input to moderation checks and conversation persistence; switched to model_dump(exclude_none=True) to filter None fields from response payloads.
Type and Responses Utilities
src/utils/responses.py, src/utils/types.py, src/utils/query.py
Introduced public types InputTool, InputToolFileSearch, InputToolMCP; replaced dict-based tool representations with structured objects throughout; expanded ResponsesApiParams with conversation, include, max_infer_iters, max_tool_calls, metadata, parallel_tool_calls, previous_response_id, prompt, stream, temperature, text, tool_choice fields; standardized Optional typing in ~15+ function signatures; made TranscriptMetadata fields optional.
Tests
tests/unit/utils/test_responses.py
Updated test assertions to use new InputToolFileSearch and InputToolMCP types with attribute access (e.g., tool.type, tool.headers) instead of dict-style access (e.g., tool["type"], tool["headers"]).

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

Suggested reviewers

  • are-ces
  • tisnik
  • jrobertboos
🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main focus of the PR, which is refactoring tool utilities to use structured API types instead of dicts and improving type safety throughout the codebase.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

try:
moderation_result = await run_shield_moderation(
client, responses_params.input, shield_ids
client, cast(str, responses_params.input), shield_ids
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Query endpoints always use string-like inputs so it is safe to cast.

response = await client.responses.create(**responses_params.model_dump())
response = await client.responses.create(
**responses_params.model_dump(exclude_none=True)
)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Explicitly drop unset attributes (queries use reduced set of responses attributes).

return tool.get("vector_store_ids", [])
return []
if tool.type == "file_search":
vector_store_ids.extend(tool.vector_store_ids)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor functional change - solves the case when there are multiple file_search objects.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, it might be relevant in the future! 👍

tools.append(tool_def)
authorization = headers.pop("Authorization", None)
tools.append(
InputToolMCP(
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Creates the whole object at once

tool_results: list[dict[str, Any]] = Field(default_factory=list)


ResponseItem: TypeAlias = (
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved at the top of the module

ResponseInput: TypeAlias = str | list[ResponseItem]


class ResponsesApiParams(BaseModel):
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Extended model with all currently available attributes in LLS.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/app/endpoints/query.py`:
- Around line 296-305: Replace the unsafe typing cast usage by converting
structured input to a string before passing it to moderation and conversation
functions: call content_to_str(responses_params.input) and use that returned
string in run_shield_moderation(...) and append_turn_to_conversation(...).
Update the calls where cast(str, responses_params.input) appears (refer to
responses_params.input, run_shield_moderation, and append_turn_to_conversation)
so they pass the content_to_str(...) result instead, ensuring you import
content_to_str from src.utils.types if not already imported.

In `@src/app/endpoints/streaming_query.py`:
- Around line 287-299: Replace the unsafe type-only cast of
responses_params.input with a runtime coercion using content_to_str from
utils.types: import content_to_str and use
content_to_str(responses_params.input) when calling run_shield_moderation and
append_turn_to_conversation (locations involving run_shield_moderation,
moderation_result, append_turn_to_conversation and responses_params.input) so
both string and list[ResponseItem] inputs are handled correctly at runtime;
mirror the same change in the similar patterns found in
src/app/endpoints/query.py.

In `@src/utils/responses.py`:
- Around line 298-299: The docstring for the parameter "tools" still mentions
"typed InputTool or dict" which is stale; update the docstring to state that
"tools" is a sequence/list of typed InputTool instances (or the exact type used
in this module) and remove any mention of dict so it reflects the typed-only API
(referencing the "tools" parameter and the InputTool type in the docstring).

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 9b21300 and 0965026.

📒 Files selected for processing (6)
  • src/app/endpoints/query.py
  • src/app/endpoints/streaming_query.py
  • src/utils/query.py
  • src/utils/responses.py
  • src/utils/types.py
  • tests/unit/utils/test_responses.py

Copy link
Contributor

@tisnik tisnik left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, indeed, now it's more readable

@tisnik tisnik merged commit ca21850 into lightspeed-core:main Feb 26, 2026
23 of 24 checks passed
tool.server_url: tool.headers
for tool in tools
if tool.get("type") == "mcp" and tool.get("headers") and tool.get("server_url")
if tool.type == "mcp" and tool.headers
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have assurance that server_url won't be e.g. "" or None (I think InputToolMCP(server_label="x", server_url="") would be valid but we don't want this)? If not, better bring back and tool.server_url.


Parameters:
tools: The prepared tools list from ResponsesApiParams.
tools: The prepared tools list (typed InputTool or dict).
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given the change you've made, this should be more accurate?
The prepared tools list of InputTool objects.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants